rm(list=ls())
library(fields); library(maps); library(ncdf4);library(abind)
library(geosphere);library(mapdata)
setwd('/Volumes/Disk2/JJA_ozone_MAM/Harvard_Dataverse')


#----- load functions -----------------------
source("./Function/get_geo.R")
source("./Function/get_met.R")
source('./Function/read_detrend.R')
source('./Function/eof.mca.R')
source('./Function/cov4gappy.R')
source('./Function/Henderson.R')
source('./Function/filled.contour3.R')

surf_T=read_SST('sst.mnmean.nc',3,5,xlim=c(100,360),ylim=c(-20,80))
lon.frame=surf_T$lon
lat.frame=surf_T$lat
yearly_O3=surf_T$data
xx=1948:2013
time.ind=(xx>=1948 & xx<=2013)


#--- area ------
lon=lon.frame;lat=lat.frame
SST_area=array(NA,c(length(lon),length(lat)))
delta1=diff(lon)[1]/2;delta2=diff(lat)[1]/2
lon0=267.5-360
for (j in 1:length(lat)){
	 x1=lon0;x2=lat[j]
     p <- rbind(c(x1+delta1,x2+delta2),c(x1+delta1,x2-delta2),c(x1-delta1,x2-delta2),c(x1-delta1,x2+ delta2)) 	
     SST_area[,j]=areaPolygon(p)
}
plot.field(SST_area,lon,lat)
sqrt(sum(SST_area)/(4*pi))


#========= Calculate the SSTs =================
ind1=(lon.frame>=(-90+360) & lon.frame<=(-30+360))
ind2=(lat.frame>=5 & lat.frame<=22)
y=array(NA,dim(yearly_O3)[3])
for(k in 1:length(y)){
	y[k]=sum(yearly_O3[ind1,ind2,k]* SST_area[ind1,ind2],na.rm=TRUE)/sum(SST_area[ind1,ind2])
}
TA.SST=mov.detrend(y[time.ind],len=7)

ind1=(lon.frame>=(-160+360) & lon.frame<=(-130+360))
ind2=(lat.frame>=32 & lat.frame<=50)
y=array(NA,dim(yearly_O3)[3])
for(k in 1:length(y)){
	y[k]=sum(yearly_O3[ind1,ind2,k]* SST_area[ind1,ind2])/sum(SST_area[ind1,ind2])
}
Pacific.SST=-mov.detrend(y[time.ind],len=7)


######## Plot begins ###############
col.gen = colorRampPalette(c("#0563FA", "white","#ff4d4d"), space="rgb")
levels=seq(-0.7, 0.7, length=61)
mid=(levels[1:(length(levels)-1)]+levels[2:(length(levels))])/2
cols=col.gen(length(levels)-1)
if(min(abs(mid))>0){
	ind1=which(mid<=0);ind2=which(mid>=0)
	ind=abind(Reduce(intersect, list(ind1+1, ind2)), Reduce(intersect, list(ind1, ind2-1)))
	ind=sort(unique(ind))
  	 cols[ind]="white"
}	

xlimit=c(180,360)
ylimit=c(0,70)

mai=c(0.05, 0.05, 0.05, 0.05)
mgp=c(1.2, 0.4, 0)
tcl=-0.2
ps=12

m <- rbind(c(0,1, 0.90,1),
c(0.05,0.5,0.70,0.90),c(0.5,0.95,0.70,0.90),
c(0.05,0.5,0.50,0.70),c(0.5,0.95,0.50,0.70),
c(0.05,0.5,0.30,0.50),c(0.5,0.95,0.30,0.50),
c(0.05,0.5,0.10,0.30),c(0.5,0.95,0.10,0.30),
c(0.1,0.9,0.02,0.42)
)
close.screen(all.screens = TRUE)
split.screen(m)
screen(1)
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
text(x=0.5,y=0.51,"Correlations of MAM SSTs with JJA met fields", font=1, cex=1.1)
text(x=0.25,y=0.08,"Tropical Atlantic SST", font=1, cex=1.0)
text(x=0.75,y=0.08,"Northeast Pacific SST", font=1, cex=1.0)
#=============== Part 1 ==================================
month=7
met=read_SST('sst.mnmean.nc',month-1,month+1,xlim= xlimit,ylim= ylimit)
lon=met$lon;lat=met$lat
month_SST=array(NA,dim(met$data[,, time.ind]))
for (i in 1:dim(month_SST)[1]){
	for (j in 1:dim(month_SST)[2]){
		month_SST[i,j,]=mov.detrend(met$data[i,j, time.ind],len=7)
	}
}

screen(2)
ap=find.cor2(month_SST,TA.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(a)")
mtext("SST",side=2,font=1,cex=1.05,line=0.1)
loc=c(360-90,360-30,5,22)
arrows(x0=loc[1],y0=loc[3],x1=loc[2],code=0,col=4,lwd=1.5)
arrows(x0=loc[1],y0=loc[4],x1=loc[2],code=0,col=4,lwd=1.5)
arrows(x0=loc[1],y0=loc[3],y1=loc[4],code=0,col=4,lwd=1.5)
arrows(x0=loc[2],y0=loc[3],y1=loc[4],code=0,col=4,lwd=1.5)

screen(3)
ap=find.cor2(month_SST,Pacific.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(b)")
loc=c(360-160,360-130,32,50)
arrows(x0=loc[1],y0=loc[3],x1=loc[2],code=0,col=4,lwd=1.5)
arrows(x0=loc[1],y0=loc[4],x1=loc[2],code=0,col=4,lwd=1.5)
arrows(x0=loc[1],y0=loc[3],y1=loc[4],code=0,col=4,lwd=1.5)
arrows(x0=loc[2],y0=loc[3],y1=loc[4],code=0,col=4,lwd=1.5)

#=============== Part 2 ==================================
met=read_surface('air.mon.mean.nc',month-1,month+1,xlim= xlimit,ylim= ylimit)
lon=met$lon
lat=met$lat
month_SST=array(NA,dim(met$data[,, time.ind]))
for (i in 1:dim(month_SST)[1]){
	for (j in 1:dim(month_SST)[2]){
		month_SST[i,j,]=mov.detrend(met$data[i,j, time.ind],len=7)
	}
}
screen(4)
ap=find.cor2(month_SST,TA.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(c)")
mtext("Surface T",side=2,font=1,cex=1.05,line=0.1)
screen(5)
ap=find.cor2(month_SST,Pacific.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(d)")

#=============== Part 3 ==================================
met=read_surface('slp.mon.mean.nc',month-1,month+1,xlim= xlimit,ylim= ylimit)
lon=met$lon
lat=met$lat
month_SST=array(NA,dim(met$data[,, time.ind]))
for (i in 1:dim(month_SST)[1]){
	for (j in 1:dim(month_SST)[2]){
		month_SST[i,j,]=mov.detrend(met$data[i,j, time.ind],len=7)
	}
}
screen(6)
ap=find.cor2(month_SST,TA.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(e)")
mtext("SLP",side=2,font=1,cex=1.05,line=0.1)
screen(7)
ap=find.cor2(month_SST,Pacific.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(f)")

#=============== Part 4 ==================================
met=read_US('hgt.mon.mean.nc',300,month-1,month+1,xlim= xlimit,ylim= ylimit)
lon=met$lon
lat=met$lat
month_SST=array(NA,dim(met$data[,, time.ind]))
for (i in 1:dim(month_SST)[1]){
	for (j in 1:dim(month_SST)[2]){
		month_SST[i,j,]=mov.detrend(met$data[i,j, time.ind],len=7)
	}
}
screen(8)
ap=find.cor2(month_SST,TA.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(g)")
mtext("300 hPa gph",side=2,font=1,cex=1.05,line=0.1)
screen(9)
ap=find.cor2(month_SST,Pacific.SST,p_value=1)
ap[ap>=0.7]=0.7;ap[ap<=-0.7]=-0.7
par(mai=mai, mgp=mgp, tcl=tcl, ps=ps)
filled.contour3(lon,lat,ap, levels=levels,  plot.axes=FALSE, col=cols)
map("world2", add = T)	
cr=find.Rlim(0.05,sum(time.ind))
clines=contourLines(lon, lat, ap,level=c(-cr, cr))
for(k in 1:length(clines)){
	lines(clines[[k]]$x,clines[[k]]$y,lwd=1,col=1,lty=2)
}
text(x=190,y=6,"(h)")

##### plot the colorbar
screen(10)
mai=c(0.1, 0.1, 0.05, 0.2)
par(mai=mai, mgp=mgp, tcl=tcl, ps=9)
image.plot(NA,legend.shrink=0.9,legend.only=TRUE,zlim=range(levels),col=cols,horizontal=TRUE,legend.width=2,axis.args = list(mgp = c(0, 0.5, 0),tcl=-0.2,padj=-1))
text("r",x=par("usr")[2]-0.13,y=par("usr")[3]+0.04,srt=0, adj = 0,xpd=TRUE,cex=1.9,font=1)

